package momo.multitree.simulation;

import java.io.FileReader;
import java.io.IOException;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Random;
import java.util.Set;
import java.util.Vector;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import momo.multitree.structure.Edge;
import momo.multitree.structure.Graph;
import momo.multitree.structure.Node;
import momo.multitree.structure.Tree;

public class SimulatedAnnealing extends TreeOptAlg
{
	static final int DEFAULT_MAXITERATION = 100;
	static final int DEFAULT_MAXTRIAL = 100;
	static final double DEFAULT_TEMPERATURE = 100d;
	static final double DEFAULT_FACTOR = 0.95d;
	
	private Log log = LogFactory.getLog(SimulatedAnnealing.class);
	

	private final long RANDOM_SEED = 894563218;
	private Random rand;
	private long seed;
	
	private int maxIteration, maxTrial;
	private double temperature, factor;
	
	public SimulatedAnnealing(Graph g, int maxIteration, int maxTrial, double temperature, double factor)
	{
		super(g);
		this.maxIteration = maxIteration;
		this.maxTrial = maxTrial;
		this.temperature = temperature;
		this.factor = factor;
		
		initBestTree();
		seed = RANDOM_SEED;
		rand = new Random(seed);
	}
	
	public SimulatedAnnealing(Graph g, Tree initBestTree, int maxIteration, int maxTrial, double temperature, double factor) {
		super(g);
		this.maxIteration = maxIteration;
		this.maxTrial = maxTrial;
		this.temperature = temperature;
		this.factor = factor;
		
		this.initBestTree = initBestTree;
		
		seed = RANDOM_SEED;
		rand = new Random(seed);
	}
	
	public SimulatedAnnealing(Graph g)
	{
		this(g, DEFAULT_MAXITERATION, DEFAULT_MAXTRIAL, DEFAULT_TEMPERATURE, DEFAULT_FACTOR);
	}
	
	public Tree optimizeTree()
	{
		return saOptTree(maxIteration, maxTrial, temperature, factor);
	}
	
//	public double getDistOfBestTree() {
//		return dist( optimizeTree() );
//	}
	
	public double score(Tree t) {
		return scoreLen(t);
	}
	
	private Tree saOptTree(int maxIteration, int maxTrial, double temperature, double factor)
	{
		Tree bestTree = initBestTree;
		Tree currTree = bestTree;

		final double initBestFitness = score(initBestTree);
//		log.info( initBestFitness );

		double curFitness = initBestFitness;

		for (int iteration = 0; iteration < maxIteration; iteration++ )
		{
			temperature = factor * temperature;
			for (int trial = 0; trial < maxTrial; trial++) 
			{
//				Tree newTree = swap(currTree);
				Tree newTree = newSwap(currTree);
				double newFitness = score(newTree);
				if ( newFitness < curFitness )
				{ // accept as better tree 
					//log.info("currFitness: "+curFitness+" new fitness: "+newFitness+" accepted!");
					bestTree = newTree;
					currTree = newTree;
					curFitness = newFitness;					
				} else if (temperature != 0) {  // accept with probability 
					double random = rand.nextDouble();
					double probability = Math.exp( (double)(curFitness-newFitness)/( temperature * initBestFitness ));
					if ( random < probability )
					{
						//log.info("currFitness: "+curFitness+" new fitness: "+newFitness+" probability: "+probability+" random: "+random+" accepted!");
						currTree = newTree;
						curFitness = newFitness;						
					}
					else { // rejected based on probability 
						//log.info("currFitness: "+curFitness+" new fitness: "+newFitness+" probability: "+probability+" random: "+random+" rejected!");
					}
				}
				else { // rejected as random walk 
					//log.info("currFitness: "+curFitness+" new fitness: "+newFitness+" probability: "+probability+" random: "+random+" rejected!");
				}
//				log.debug(""+iteration+" "+" "+temperature+" "+trial+" "+curFitness);
			} // end for: round
//			log.info(""+iteration+", " + curFitness );
		} // end for: iteration
		
		return bestTree;
	}
	
	private void initBestTree()
	{
		double e1 = score(bestLatTree);
		double e2 = score(bestCostTree);
		double e3 = score(bestStabTree);
		
		if ( e1 < e2 )
			if ( e1 < e3 )
				initBestTree = bestLatTree;
			else
				initBestTree = bestStabTree;
		else
			if ( e2 < e3 )
				initBestTree = bestCostTree;
			else
				initBestTree = bestStabTree;
	}
	
	private Tree swap(Tree tree)
	{
		Tree newTree;
		int noOfNodes = tree.getNoOfNodes();
		int noOfEdges = tree.getNoOfEdges();
		double probOfEdge = (double)noOfEdges / (double)(noOfEdges + noOfNodes);
		double prob = rand.nextDouble();
		if ( prob <= probOfEdge )
			newTree = swapBranch(tree);
		else
			newTree = swapRoot(tree);
		return newTree;
	}
	
	private void testSwap() {
		Tree curTree = initBestTree;
		for(int i=0; i<10; i++) {
			curTree = newSwap(curTree);
		}
	}
	
	private Tree newSwap(Tree tree) {
		Set<Node> targetNodes = new LinkedHashSet<Node> ( g.getNodes() );
		targetNodes.removeAll( g.getLeafNodes() );
		targetNodes.remove( tree.getRoot() );
		
		if ( targetNodes.size() == 0 )
			return tree;
		
		Map<Node, Double> map = new LinkedHashMap<Node, Double>();
		double p = 1.0;
		map.put(null, p);
		for(Node n : targetNodes) {
			double curP = 1.0/(1.0 + tree.getDescendants(n).size());
			p += curP;
			map.put(n, curP);
		}
		p = 1.0/p;
		
		double r = rand.nextDouble();
		double currCum = 0;
		for (Node n: map.keySet()) {
			currCum += map.get(n) * p;
			assert ( 1.0 - currCum > -1e-5 ) : "value of currCom exceeds 1.0: "+currCum;
			if ( r < currCum ) {
				if (n == null) {
					// root swap
					Tree newTree = swapRoot(tree);
//					newTree.checkParentMap();
					return newTree;
				}
				else { // uplink swap
					Edge e = tree.getEdge(n, tree.getParent(n) );
					
					assert e!=null;
					Tree newTree = swapBranch(tree, e);
					return newTree;
				}
			}
		}
		assert false;		
		return null;
	}
	
	private Tree swapRoot(Tree tree)
	{
		Tree newTree = tree.clone();
		Node[] nodes = newTree.getNodes().toArray(new Node[0]);
		int pos = getRandNo(0, nodes.length-1);
		newTree.changeRoot(nodes[pos]);
		return newTree;
	}
	
	private Tree swapBranch(Tree tree)
	{
//		System.out.println(tree.toString()+"");
		
		Tree newTree = tree.clone();
		boolean swap = false;
		Vector<Edge> toRemove = getNonLeafTreeEdges(newTree);
		
		if ( toRemove.size() == 0 ) 
			return newTree;
		
		int posToRemove = getRandNo(0, toRemove.size()-1);
		Edge removalEdge = toRemove.remove(posToRemove);
		Vector<Edge> toSwap = getRemainingEdges(newTree);
		
//		System.out.println("Removal Edge: " + removalEdge);
		newTree.remove(removalEdge);
		
		while (!swap && toSwap.size() > 0 )
		{
			int posToSwap = getRandNo(0, toSwap.size()-1);
			Edge swapEdge = toSwap.remove(posToSwap);
			
			if ( !newTree.hasCycle(swapEdge) )
			{
				newTree.add(swapEdge);
				swap = true;
			}
		}
		
		if ( !swap )
			newTree.add(removalEdge);
		
//		System.out.println("\n\n");
		return newTree;
	}
	
	private Tree swapBranch(Tree tree, Edge removalEdge)
	{
//		System.out.println(tree.toString()+"");
		
		Tree newTree = tree.clone();
//		newTree.checkParentMap();
		
		boolean swap = false;
		Vector<Edge> toSwap = getRemainingEdges(newTree);
				
		boolean hasRemoved = newTree.remove(removalEdge);
		
		while (!swap && toSwap.size() > 0 )
		{
			int posToSwap = getRandNo(0, toSwap.size()-1);
			Edge swapEdge = toSwap.remove(posToSwap);
			
			if ( !newTree.hasCycle(swapEdge) )
			{
				newTree.add(swapEdge);
//				newTree.checkParentMap();
				swap = true;
			}
		}
		
		if ( !swap ) {
			newTree.add(removalEdge);
//			newTree.checkParentMap();
		}
		
//		System.out.println("\n\n");
		return newTree;
	}
	
	private Vector<Edge> getRemainingEdges(Tree tree)
	{
		Vector<Edge> vec = new Vector<Edge>();
		Set<Edge> graphSet = g.getNonLeafEdges();
		for(Edge e: graphSet)
			if ( !tree.contains(e) )
				vec.add(e);
		return vec;
	}
	
	private Vector<Edge> getNonLeafTreeEdges(Tree tree)
	{
		Vector<Edge> vec = new Vector<Edge>();
//		Set<Edge> treeSet = tree.getEdges();
		Set<Edge> treeSet = new LinkedHashSet<Edge>(tree.getEdges());
		treeSet.removeAll(g.getLeafEdges());
		vec.addAll(treeSet);
		return vec;
	}
	
	public int getRandNo(int low, int high)
	{
		return rand.nextInt(high-low+1) + low;
	}
	
	public static void main(String args[]) {
		try {
			String filename = "50_node_symmetric_noleaf_dataset_55.txt";
			Graph g = new Graph(new FileReader("dataset/"+filename));
			SimulatedAnnealing sim = new SimulatedAnnealing(g, 1000, 10, 0.01, 0.95);
			sim.optimizeTree();
			
		}catch (Exception e) {
			e.printStackTrace();
		}
	}
			
}//end of class SimulatedAnnealing
